from langchain.tools import tool
import os
from langchain.chat_models import init_chat_model
from dotenv import load_dotenv
from langchain_classic.agents import create_tool_calling_agent, AgentExecutor
from langchain_core.prompts import ChatPromptTemplate

# 加载.env文件

load_dotenv()

# 第一步：定义加法、减法工具
# 通过@tool装饰器分别定义两个工具，关键是在文档字符串中明确区分功能，让模型能精准识别适用场景。

# 加法工具
@tool
def add_tool(num1: float, num2: float) -> str:
    """
    用于执行两个数字的加法运算。
    当用户需要计算“A加B”“A加上B”“A与B的和”“A和B相加”等场景时使用。

    Args:
        num1: 加法运算中的第一个数字（整数或小数均可）
        num2: 加法运算中的第二个数字（整数或小数均可）
    """
    result = num1 + num2
    return f"加法运算结果：{num1} + {num2} = {result}"

# 减法工具
@tool
def subtract_tool(num1: float, num2: float) -> str:
    """
    用于执行两个数字的减法运算。
    当用户需要计算“A减B”“A减去B”“A与B的差”“A比B多多少”“A减去B等于多少”等场景时使用。

    Args:
        num1: 减法运算中的被减数（整数或小数均可）
        num2: 减法运算中的减数（整数或小数均可）
    """
    result = num1 - num2
    return f"减法运算结果：{num1} - {num2} = {result}"


# 第二步：搭建 Agent 让模型驱动工具选择
# 将两个工具传入 Agent，结合大语言模型（LLM），让模型根据用户输入的问题，自主判断该调用加法还是减法工具。
# 1. 初始化模型
model = init_chat_model(
    model='deepseek-reasoner', # deepseek-chat表示调用DeepSeek-v3模型，deepseek-reasoner表示调用DeepSeek-R1模型，
    model_provider='deepseek',# 模型提供商写deepseek
    # api_key="sk-476b894aa2d74474b77d4dfa91633f72", #你注册的deepseek api_key
    api_key=os.environ.get("DEEPSEEK_API_KEY"),
)

# 2. 定义Agent的提示词（引导模型正确使用工具）
prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个擅长数学计算的助手，需要根据用户的计算需求，选择合适的工具进行运算。"
               "如果是加法需求，调用add_tool；如果是减法需求，调用subtract_tool。"
               "严格按照工具要求的参数格式传参，不要随意增减参数。"),
    ("human", "{input}"),  # 用户输入的问题
    ("placeholder", "{agent_scratchpad}")  # Agent思考和工具调用的临时区域
])


# 3. 创建工具调用Agent（传入模型、工具列表、提示词）
tools = [add_tool, subtract_tool]
agent = create_tool_calling_agent(model, tools, prompt)


# 4. 创建Agent执行器（负责运行Agent并处理工具调用流程）
agent_executor = AgentExecutor(
    agent=agent,
    tools=tools,
    verbose=True,  # 开启verbose模式，可查看模型思考和工具调用过程
    handle_parsing_errors=True  # 自动处理参数解析错误
)

# 第三步：测试模型的工具选择能力
# 通过agent_executor.invoke()传入不同的计算需求，观察模型是否能正确选择加法或减法工具。

# 输入加法问题
result1 = agent_executor.invoke({"input": "计算一下15.6加上23.4等于多少？"})
print("最终回答1：", result1["output"])

# 测试案例 2：减法需求
# 输入减法问题
result2 = agent_executor.invoke({"input": "50减去12.8等于多少？"})
print("最终回答2：", result2["output"])